#include "structure_amova.h"

StructureAMOVA::StructureAMOVA( DataSet *theData ) : AnalysisBase( theData ) {
}

bool StructureAMOVA::run() {
    if( !H )
        return false;

    _results->appendHeader("AMOVA Results");


    createDistanceMatrix();
    if( !D ) return false;
    createCovarianceMatrix();
    if( !D ) return false;

    sst = MatrixTrace(D);

    gsl_matrix *R = gsl_matrix_alloc( D->size1, D->size2 );
    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans,
                   1.0, H, D,
                   0.0, R);
    ssw = MatrixTrace(R);
    ssa = sst-ssw;

    return true;
}


void StructureAMOVA::setPopulationPartition(QString stratum) {
    createHMatrix( stratum );
}

gsl_matrix* StructureAMOVA::getHMatrix( QString stratum ) {
    createHMatrix( stratum );
    return H;
}

gsl_matrix* StructureAMOVA::getDistanceMatrix() {
    createDistanceMatrix();
    return D;
}

gsl_matrix* StructureAMOVA::getCovarianceMatrix() {
    createDistanceMatrix();
    createCovarianceMatrix();
    return D;
}


void StructureAMOVA::createDistanceMatrix() {
    unsigned N = (uint)_theData->numInds();
    D = gsl_matrix_calloc(N,N);
    QStringList locusNames = _theData->getIndividual(0)->locusNames();

    for(uint i=0;i<N;i++){
        Individual *ind1 = _theData->getIndividual(i);
        for(uint j=i+1;j<N;j++){
            Individual *ind2 = _theData->getIndividual(j);
            qreal genDist = 0.0;
            foreach(QString locus, locusNames)
                genDist += getLocusAMOVADistance( ind1->getLocus(locus),
                                                  ind2->getLocus(locus),
                                                  _theData->getTotalFrequencies()->getLocusFrequency(locus ) );
            gsl_matrix_set(D,i,j,genDist);
            gsl_matrix_set(D,j,i,genDist);
        }
    }
}

void StructureAMOVA::createCovarianceMatrix() {
    if( !D )
        createDistanceMatrix();
    unsigned N = D->size1;
    gsl_matrix *C = gsl_matrix_calloc(N,N);
    gsl_vector *v = gsl_vector_calloc(N);
    qreal tot = 0.0;

    // create mean vector
    for(uint i=0;i<N;i++){
        for(uint j=i+1;j<N;j++){
            gsl_vector_set( v, i, gsl_matrix_get(D,i,j)+gsl_vector_get(v,i));
            gsl_vector_set( v, j, gsl_matrix_get(D,j,i)+gsl_vector_get(v,j));
            tot += gsl_matrix_get(D,i,j) + gsl_matrix_get(D,j,i);
        }
    }
    gsl_vector_scale(v,1.0/N);
    tot /= (N*N);

    // create covariance matrix
    for(uint i=0;i<N;i++){
        for(uint j=i+1;j<N;j++){
            qreal val = gsl_matrix_get(D,i,j) - gsl_vector_get(v,i) - gsl_vector_get(v,j) + tot;
            gsl_matrix_set(C,i,j,val);
            gsl_matrix_set(C,j,i,val);
        }
    }
    gsl_matrix_memcpy(D,C);
    gsl_matrix_free( C );
}

void StructureAMOVA::createHMatrix(QString stratum) {
    _theData->partitionByStratum(stratum);
    QStringList strataNames = _theData->getStrataNames();
    uint N = _theData->numInds();
    uint K = _theData->numPops();
    gsl_vector *ni = gsl_vector_calloc(K);
    for(uint i=0;i<K;i++)
        gsl_vector_set(ni,i,_theData->popSize(strataNames[i]));

    H = gsl_matrix_calloc( N, N );
    for(uint i=0;i<N;i++){
        QString strat1 = _theData->getIndividual(i)->getStrata(stratum);
        for(uint j=i+1;j<N;j++){
            if( _theData->getIndividual(j)->getStrata(stratum) == strat1 ){
                int idx = strataNames.indexOf(strat1);
                qreal val = 1.0 / gsl_vector_get(ni,idx);
                gsl_matrix_set(H,i,j,val);
                gsl_matrix_set(H,j,i,val);
            }
        }
    }
}
